import matplotlib.pyplot as plt


# ============================================= 画图相关 ========================================
# 保存图片: 2幅 * 1行
def save_images_2x1(path_img, data1, data1_str, data2, data2_str, str_label):
    # 创建一个新的图像
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    # 在左上方绘制 list_data1[i]
    im1 = ax1.imshow(data1, aspect='auto')
    ax1.set_title(data1_str)
    ax1.axis('equal')
    ax1.axis('off')
    # plt.colorbar(im1, ax=ax1)     # 注释掉能关闭图例
    # 在中上方绘制 list_data2[i]
    im2 = ax2.imshow(data2, aspect='auto')
    ax2.set_title(data2_str)
    ax2.axis('equal')
    ax2.axis('off')
    # plt.colorbar(im2, ax=ax2)
    # 在下方放置文本
    fig.text(0.5, 0.01, str_label, ha='center', fontsize=12)
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, bottom=0.05)
    # 显示图像
    # plt.show()
    # 保存图像
    plt.savefig(path_img)
    # 关闭当前图像，释放内存
    plt.close(fig)


# 画图: 2幅 * 1行
def plot_images_2x1(data1, data1_str, data2, data2_str, str_label):
    # 创建一个新的图像
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    # 在左上方绘制 list_data1[i]
    im1 = ax1.imshow(data1, aspect='auto')
    ax1.set_title(data1_str)
    ax1.axis('equal')
    ax1.axis('off')
    # plt.colorbar(im1, ax=ax1)     # 注释掉能关闭图例
    # 在中上方绘制 list_data2[i]
    im2 = ax2.imshow(data2, aspect='auto')
    ax2.set_title(data2_str)
    ax2.axis('equal')
    ax2.axis('off')
    # plt.colorbar(im2, ax=ax2)
    # 在下方放置文本
    fig.text(0.5, 0.01, str_label, ha='center', fontsize=12)
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, bottom=0.05)
    # 保存图像
    plt.show()


# 画图: 2幅 * 2行 + 1行文字
def plot_images_2x2(data1, data2, data3, data4, img_text):
    # 创建一个新的图像，包含2行2列的子图
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    # 在第一行绘制两张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(f"phaseBit(origin)")
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(f"phaseBit(decoder)")
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    # 在第二行绘制两张图
    im3 = axes[1, 0].imshow(data3, aspect='auto')
    axes[1, 0].set_title(f"pattern(origin)")
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im4 = axes[1, 1].imshow(data4, aspect='auto')
    axes[1, 1].set_title(f"pattern(decoder))")
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    plt.show()


# 保存图: 2幅 * 2行 + 1行文字
def save_images_2x2(path_img,
                    data1, text1, data2, text2, data3, text3, data4, text4,
                    img_text):
    # 创建一个新的图像，包含2行2列的子图
    fig, axes = plt.subplots(2, 2, figsize=(18, 9))
    # 在第一行绘制两张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(text1)
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(text2)
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    # 在第二行绘制两张图
    im3 = axes[1, 0].imshow(data3, aspect='auto')
    axes[1, 0].set_title(text3)
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im4 = axes[1, 1].imshow(data4, aspect='auto')
    axes[1, 1].set_title(text4)
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    # plt.show()
    # 保存图像
    plt.savefig(path_img)
    # 关闭当前图像，释放内存
    plt.close(fig)


# 画图: 3幅 * 2行 + 1行文字
def plot_images_3x2(data1, data2, data3, data4, data5, data6, img_text):
    # 创建一个新的图像，包含2行3列的子图
    fig, axes = plt.subplots(2, 3, figsize=(18, 9))
    # 在第一行绘制三张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(f"phaseBit(t)")
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(f"phaseBit-NN(t-1)")
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    im3 = axes[0, 2].imshow(data3, aspect='auto')
    axes[0, 2].set_title(f"phaseBit-NN(t)")
    axes[0, 2].axis('equal')
    axes[0, 2].axis('off')
    # 在第二行绘制三张图
    im4 = axes[1, 0].imshow(data4, aspect='auto')
    axes[1, 0].set_title(f"pattern")
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im5 = axes[1, 1].imshow(data5, aspect='auto')
    axes[1, 1].set_title(f"diff(phaseBit(t), phaseBit-NN(t))")
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    im6 = axes[1, 2].imshow(data6, aspect='auto')
    axes[1, 2].set_title(f"pattern-NN")
    axes[1, 2].axis('equal')
    axes[1, 2].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    plt.show()


# 保存图: 3幅 * 2行 + 1行文字
def save_images_3x2(path_img,
                    data1, text1, data2, text2, data3, text3,
                    data4, text4, data5, text5, data6, text6,
                    img_text):
    # 创建一个新的图像，包含2行3列的子图
    fig, axes = plt.subplots(2, 3, figsize=(18, 9))
    # 在第一行绘制三张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(text1)
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(text2)
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    im3 = axes[0, 2].imshow(data3, aspect='auto')
    axes[0, 2].set_title(text3)
    axes[0, 2].axis('equal')
    axes[0, 2].axis('off')
    # 在第二行绘制三张图
    im4 = axes[1, 0].imshow(data4, aspect='auto')
    axes[1, 0].set_title(text4)
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im5 = axes[1, 1].imshow(data5, aspect='auto')
    axes[1, 1].set_title(text5)
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    im6 = axes[1, 2].imshow(data6, aspect='auto')
    axes[1, 2].set_title(text6)
    axes[1, 2].axis('equal')
    axes[1, 2].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    # plt.show()
    # 保存图像
    plt.savefig(path_img)
    # 关闭当前图像，释放内存
    plt.close(fig)
